package dataframe

import org.apache.spark.SparkConf
import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.functions.udf

object DataFrame_UDFTest {
  def main(args: Array[String]): Unit = {
    val conf = new SparkConf()
    conf.setMaster("local[*]")
    conf.setAppName("DataFrame_UDFTest")

    val spark: SparkSession = SparkSession
      .builder()
      .config(conf)
      .getOrCreate()
    import spark.implicits._

    // 在Scala中一个简单的UDF，将数字等级转换为考查等级
    // 创建学生成绩DataFrame
    val studentDF = Seq(
      Student("张三", 85),
      Student("李四", 90),
      Student("王老五", 55)
    ).toDF()

    // 创建一个函数(普通的Scala函数)将成绩转换到考察等级
    studentDF.createOrReplaceTempView("tb_student")

    spark.sql("select * from tb_student").show()

    // 注册为一个UDF（在DSL API中使用时的注册方法）
    val convertGradeUDF = udf(convertGrade(_: Int): String)

    // 使用该UDF将成绩转换为字母等级
    studentDF.select($"name", $"score", convertGradeUDF($"score").as("grade")).show()


    // 在SQL查询中使用UDF时，注册过程略有不同
    println("*"*100)
    spark.udf.register("convertGrade_SQL",convertGrade(_:Int):String)
    spark.sql("select name,score,convertGrade_SQL(score) as grade from tb_student").show()


    spark.stop()
  }

  // 创建一个函数(普通的Scala函数)将成绩转换到考察等级
  def convertGrade(score: Int): String = {
    score match {
      case `score` if score > 100 => "作弊"
      case `score` if score >= 90 => "优秀"
      case `score` if score >= 80 => "良好"
      case `score` if score >= 70 => "中等"
      case `score` if score >= 60 => "及格"
      case _ => "不及格"
    }
  }

  case class Student(name: String, score: Int)

}
